from torchvision.datasets import ImageFolder
from torchvision.transforms import Compose, ToTensor, Normalize, Resize, RandomResizedCrop

from data import Data


class ImageNet(Data):
    def __init__(self, n_clients, batch_size, alpha=1000, path="/dataset/ILSVRC/Data/CLS-LOC", flag=True):
        super().__init__()
        transform = Compose([
            Resize(224),
            RandomResizedCrop(224),
            ToTensor(),
            Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225])
        ])
        self.dataset = ImageFolder(
            root=path + "/train",
            transform=transform)
        self.validate = ImageFolder(
            root=path + "/val",
            transform=transform)
        self.trainLoader, self.client_nums, self.total = \
            self.train_loader(alpha, n_clients, batch_size, flag)
        self.validationLoader = self.validate_loader(batch_size)

    def __str__(self):
        return "ImageNet"
